import os
from jinja2 import Environment, FileSystemLoader, select_autoescape
from langchain_core.prompts import PromptTemplate

from ..prompts.common import prompt_context

jinja_env = Environment(
    loader=FileSystemLoader(os.path.join(os.path.dirname(__file__), "tpl")),
    autoescape=select_autoescape(),
    trim_blocks=True,
    lstrip_blocks=True,
)
os_env = Environment(
    loader=FileSystemLoader("/"),  # 使用根目录作为基础路径
    autoescape=select_autoescape(),
    trim_blocks=True,
    lstrip_blocks=True,
)


def get_prompt_template(prompt_name: str, **kwargs) -> str:
    """
    Load and return a prompt template using Jinja2.

    Args:
        prompt_name: Name of the prompt template file (without .md extension)

    Returns:
        The template string with proper variable substitution syntax
    """
    # 获取环境变量中的模板路径
    agent_tpl_path = os.getenv('AGENT_TPL_PATH')

    try:
        # 优先检查环境变量指定的路径
        if agent_tpl_path:
            custom_template_path = os.path.join(agent_tpl_path, f"{prompt_name}.md")
            if os.path.exists(custom_template_path):
                template = os_env.get_template(custom_template_path)
                return template.render(**kwargs)

        template = jinja_env.get_template(f"{prompt_name}.md")
        return template.render(**kwargs)
    except Exception as e:
        raise ValueError(f"Error loading template {prompt_name}: {e}")


def get_prompt_context(context: str) -> str:
    return f"""
<context>
{PromptTemplate.from_template(template=prompt_context).format(context=context)}
</context>
"""
